import numpy as np
import matplotlib.pyplot as plt
import Ising_Metropolis
import wolff
import MCMC_meachinlearning
from Ising_Metropolis import nn
from numpy.random import rand

# 将模型的初始变量设计为 J_1 = 1 J_2 = -0.01 J_3 = -0.01
def flipping(grid, beta):
    N = len(grid)
    test_j = [1,0.01,0.01]
    for i in range(N):
        for j in range(N):
            a = np.random.randint(0, N)
            b = np.random.randint(0, N)
            s = grid[a][b]
            E = 0
            count = 1
            for k in test_j:
                E += k*np.array(nn(grid,(a,b),count,N)).sum()
                count += 1
            cost = 2 * s * E
            if cost < 0:
                s *= -1
            elif rand() < np.exp(-cost * beta):
                s *= -1
            grid[a][b] = s
    return grid

def metropolis(eqsteps, mcsteps, N):
    T = np.linspace(0.1, 5, 10)
    n1 = 1 / (N * N * mcsteps)
    M = []
    j = 0
    for t in T:
        m = 0
        config = Ising_Metropolis.init_state(N)
        # 先进行平衡
        for i in range(eqsteps):
            flipping(config, 1 / t)
        # 之后进行计算
        for i in range(mcsteps):
            flipping(config, 1 / t)
            m += abs(np.sum(config))
            j += 1
            if j%10 == 0:
                print("已完成第%d步模拟" % j)
        M.append(m * n1)
    sp = f.add_subplot(2,2,1)
    plt.xlabel('T')
    plt.ylabel('m')
    plt.title('metropolis')
    plt.plot(T, abs(np.array(M)), "-or")
    sp = f.add_subplot(2, 2, 3)
    plt.xlabel('T')
    plt.ylabel('m')
    plt.plot(T, abs(np.array(M)), "-or")

def wolff_1(grid,location:tuple, beta):
    test_j = [1,0.01,0.01]
    j = np.array(test_j)
    p = 1 - np.exp(-2*beta*j)
    stack = []
    x = location[0]
    y = location[1]
    s = grid[x][y]
    cluster = []
    stack.append([s,(x,y)])
    cluster.append((x,y))
    while stack:
        # first 1nn
        temp = stack.pop()
        #  依次连接最近邻格点,2近邻格点,3近邻格点
        for i in range(3):
            neigh = wolff.nn(grid,temp[1],i+1,len(grid))
            for em in neigh:
                if em[1] not in cluster and em[0] == s and rand() < p[i]:
                    stack.append(em)
                    cluster.append(em[1])
    for i in cluster:
        x = i[0]
        y = i[1]
        grid[x][y] = -1*grid[x][y]
    return grid

# 算法的应用
def wolff_flip(grid, beta, J=1, K=0.2):
    # 在有效模型中,我们要对wolff的每一次翻转进行有效性判断
    test_j = [1,0.01,0.01]
    steps = 1
    for i in range(steps):
        Eeff = [] #存放通过有效模型计算的哈密顿量
        temp = grid.copy() #保存未翻转的构型
        nn_a,E_A = MCMC_meachinlearning.heff(temp) #计算构型的近邻关系以及哈密顿量
        Eeff.append(nn_a) #存放近邻关系以便于 使用有效模型计算哈密顿量

        # wollf算法
        a = np.random.randint(0,len(grid))
        b = np.random.randint(0,len(grid))
        wolff_1(grid,(a,b),beta)

        #计算反转后的信息
        nn_b,E_B = MCMC_meachinlearning.heff(grid)
        Eeff.append(nn_b)
        Eeff_A = -1*test_j[0]*Eeff[0][0]+-1*test_j[1]*Eeff[0][1]+-1*test_j[2]*Eeff[0][2]
        Eeff_B = -1*test_j[0]*Eeff[1][0]+-1*test_j[1]*Eeff[1][1]+-1*test_j[2]*Eeff[1][2]
        if rand() > np.exp(-1*beta*((E_B-Eeff_B)-(E_A-Eeff_A))):
            grid = temp
    return grid

def wolff_test(esteps,mcsteps,N):
    T = np.linspace(0.1, 5, 10)
    n1 = 1 / (N*N*mcsteps)
    M = []
    j = 0
    for t in T:
        m = 0
        config = Ising_Metropolis.init_state(N)
        for i in range(esteps):
            wolff_flip(config,1/t)
        for i in range(mcsteps):
            wolff_flip(config,1/t)
            m += abs(np.sum(config))
            j += 1
            if j%10 == 0:
                print('已完成第%d步模拟'%j)
        M.append(m*n1)
    sp = f.add_subplot(2,2,2)
    plt.xlabel('T')
    plt.ylabel('m')
    plt.title('wollf')
    plt.plot(T, abs(np.array(M)), "-ob")
    sp = f.add_subplot(2, 2, 3)
    plt.xlabel('T')
    plt.ylabel('m')
    plt.title('metropolis compare with wollf')
    plt.plot(T, abs(np.array(M)), "-ob")
    plt.show()


if __name__ == '__main__':
    f = plt.figure(figsize=(18,10))
    metropolis(1000,1000,8)
    wolff_test(1000,1000,8)
